Run fused-MoE expert GPTQ updates on all of this process's GPUs - #3002

…3019)

### What does this PR do?

Type of change: new feature (performance)

Adds `hessian_input_dtype` to the GPTQ and `local_hessian` algorithm
configs: `"float32"` (the default, today's behavior) or `"bfloat16"`.

`update_hessian` copies every batch of activations to float32 and
accumulates `XᵀX` with a float32 GEMM, which runs on CUDA cores. With
`"bfloat16"`, a tensor-core GEMM reads the activations in bfloat16 and
writes float32 (`out_dtype=torch.float32`). For bfloat16 activations,
the products are exact in float32, so it differs from the float32 GEMM
only in how it rounds the sums. It also needs no float32 copy of each
batch, which is twice the size of the bfloat16 batch.

Against a float64 reference, the bfloat16 option's GPTQ Hessian is off
by up to 2.8e-5 of the mean diagonal at 32K tokens per batch (1.3e-5 at
4K), versus 4.7e-6 for the float32 GEMM. A bfloat16 output would be off
by 9.2e-3. GPTQ then adds 1% of the mean diagonal as damping, so the
difference does not reach the result, which matches the KL below. The
Hessian is stored in float32 either way, so the option saves only the
transient float32 copy, not Hessian memory.

- When the Hessian lives on another device, as #2945 places it by free
memory, the float32 product is copied over.
- Where CUDA GEMMs don't take `out_dtype` (CPU inputs, or an older
torch), `"bfloat16"` falls back to float32 with a one-time warning
instead of being ignored. That check is cached per device.
- It applies to every GPTQ Hessian: fused MoE experts and linear layers
alike.
- `local_hessian` builds 16x16 per-block Hessians the same way, so it
gets the same option: a batched GEMM (`bmm`) that reads bfloat16 and
writes float32. On an RTX PRO 6000 at 32K tokens it is 6.3x faster for
4,096 features (2.47 -> 0.40 ms per update) and 7.7x for 2,048, and it
no longer makes the 512 / 256 MiB float32 copy. Its error is up to
1.4e-4 of the block diagonal, versus 1.9e-5 for float32. That Hessian
only weights the error of each candidate scale in the search, and it is
never inverted.
- The field lives in a small mixin both configs share. One helper,
`bfloat16_hessian_gemm`, decides the GEMM for both and warns when it is
unavailable.
- imatrix keeps its float32 statistic: it is a per-column sum of
squares, not a GEMM, so tensor cores do not apply.

### Usage

In a recipe or config:

```yaml
algorithm:
  method: gptq
  hessian_input_dtype: bfloat16  # default: float32
```

### Testing

**GEMM alone on GB300.** Accumulating about 360 batches of 1 to 2,000
tokens each into one Hessian, steady state:

| Hessian width | float32 GEMM | bfloat16 in, float32 out |
|---|---|---|
| 2,560 (Qwen3.8-Flash-Next experts) | 96 ms | **11 ms (8.7x)** |
| 4,096 (GLM-5.3-Flash experts) | 244 ms | **15 ms (16x)** |

**Real models on 4x GB300 (oci-jhb), float32 vs bfloat16.** Layerwise
GPTQ on the routed experts at IQ2_XXS, 1,024 x 512 calibration tokens,
KL to BF16 over wikitext-2 (100 x 512 tokens, second half scored).
Measured on top of #2945 and #3002, with flash-linear-attention
installed, before this PR made the option opt-in:

| | GPTQ, all layers | Layerwise calibration | Whole job | KL to BF16
(paired delta, 95%) |
|---|---|---|---|---|
| GLM-5.3-Flash | 730 s -> **553 s (-24%)** | 16:30 -> **13:09 (-20%)**
| 40:41 -> **36:04 (-11%)** | 0.6306 -> 0.6293 (-0.0013 ± 0.0080, p
0.75) |
| Qwen3.8-Flash-Next | 624 s -> **581 s (-7%)** | 15:26 -> **14:33
(-6%)** | 37:27 -> **35:46 (-5%)** | 0.2480 -> 0.2494 (+0.0015 ± 0.0052,
p 0.58) |

The Qwen run with the bfloat16 GEMM rebuilt the GGML CUDA extension for
its new source tree, which took 142 s. The table subtracts that one-time
build, which is unrelated to this change. KL moves within run-to-run
noise on both models.

- GPU test `test_bf16_hessian_input_gives_the_float32_hessian`:
`input_dtype=torch.bfloat16` gives the float64 reference Hessian to
within 1e-4 (the diagonal is about 2), with the Hessian on the inputs'
GPU and on the CPU. Where the fast GEMM exists, an update allocates less
than a float32 copy of its batch. Without the bfloat16 path the update
allocates 16 MiB, the float32 copy, and the test fails.
- GPU test `test_gptq_hessian_input_dtype_reaches_the_hessian`: through
`mtq.quantize` with a GPTQ config, `hessian_input_dtype: bfloat16`
lowers the measured Hessian-forward peak by the float32 copy of the
batch. It fails if the option does not reach `update_hessian`.
- New GPU tests in
`tests/gpu/torch/quantization/test_local_hessian_gpu.py`.
`test_bf16_local_hessian_input_gives_the_float32_hessian` checks the
per-block Hessian against float64, and that no float32 copy is made.
`test_local_hessian_input_dtype_reaches_the_hessian` runs `mtq.quantize`
with `local_hessian`, `hessian_input_dtype: bfloat16` and `debug: true`,
and gets a Hessian that differs from the float32 one only by rounding.
Each fails without its piece.
-
`tests/gpu/torch/quantization/{test_gptq,test_layerwise_calibrate}.py`:
36 passed on 2x RTX PRO 6000. `test_gptq.py`,
`test_local_hessian_gpu.py` and `test_nvfp4_fp8_sweep_kernel.py`: 69
passed. The first version, with the bfloat16 GEMM always on, also passed
`test_gptq.py` on 4x GB300 (torch 2.13, transformers 5.16).
-
`tests/unit/torch/quantization/{test_gptq,test_ggml_gptq,test_local_hessian,test_config_validation}.py`
and `tests/unit/recipe`: 666 passed.
- New CPU tests
`test_bf16_hessian_input_on_cpu_gives_the_float32_result_and_warns_once`,
one each in `test_gptq.py` and `test_local_hessian.py`. On CPU, where
GEMMs take no `out_dtype`, `hessian_input_dtype: bfloat16` gives the
float32 result exactly and warns once over three batches. Without the
once-per-device cache they see 3 and 9 warnings and fail.
- **torch 2.8.0+cu128**, the oldest torch this package supports, on 2x
RTX PRO 6000 with transformers 5.18. `mm`, `addmm` and `bmm` all have
`out_dtype` overloads there, so every Hessian GEMM takes the fast path.
`tests/gpu/torch/quantization/{test_gptq,test_local_hessian_gpu}.py`: 24
passed, including the four `hessian_input_dtype` tests, which skip
without the fast path.
`tests/unit/torch/quantization/{test_gptq,test_local_hessian}.py`: 19
passed, on torch 2.8 and 2.12 alike. torch 2.6 has no such overloads, so
there `"bfloat16"` falls back to float32 with the warning.

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ✅ (new option; the default keeps
the float32 GEMM)
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
✅
- Did you get Claude approval on this PR?: ❌

🤖 Generated with [Claude Code](https://claude.com/claude-code)


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

<!-- coderabbit-ai-generated:v1 -->

## Summary by CodeRabbit

* **New Features**
* GPTQ and local-Hessian calibration now support an optional `bfloat16`
input mode for Hessian calculations, while retaining float32
accumulation. This can reduce calibration memory use on supported CUDA
hardware.
* The default remains `float32`. If the hardware does not support the
required operations, calibration falls back to float32 and displays a
warning.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

### What does this PR do?

Type of change: new feature (performance)

After its forward, imatrix encodes every target weight once, one weight
after another. In the default single-pass run of a model that
`device_map` spreads over several GPUs, each encode runs on its weight's
GPU. One more thing serializes them: decoding the payload reads its
shape to the host, which waits for that GPU. So all but one GPU sat
idle.

Weights that are readable as they stand are now grouped by device and
encoded by one thread per GPU, each on the GPU the weight already lives
on. Nothing moves between devices, so the codes are the same as before.

- Weights that need a materialization window are still encoded in turn
inside it: offloaded weights, FSDP2 shards and TP-sharded linears
(`requires_weight_materialization`, `is_quantized_parallel_linear`).
- The first encode runs before the workers start, so a format's CUDA
extension, which is not thread-safe to load, loads once.
- Worker threads turn off autograd themselves, since grad mode is per
thread.
- With weights on one GPU, as in a layerwise run that holds one layer at
a time, nothing changes.

This is independent of #3002, which moves GPTQ expert updates to other
GPUs. Here every weight stays where it is.

### Usage

No API change.

### Testing

- **Synthetic model, 192 bf16 4,096 x 4,096 linears over 3 RTX PRO
6000s, IQ2_XXS imatrix**, warm extension: the encode phase takes 3.2 s
on `main` and **1.2 s (2.7x)** with this PR.

- **Real models on 4x GB300 (oci-jhb)**: single-pass IQ2_XXS imatrix,
1,024 x 512 calibration tokens, the model spread over the 4 GPUs by
`device_map`. Encode phase, minus the one-time JIT build of the GGML
extension that falls in the first encode (128-144 s in each run):

  | | `main` | This PR | Whole job |
  |---|---|---|---|
| GLM-5.3-Flash (24,192 expert weights) | 334 s | **91 s (3.7x)** |
34:00 -> 30:05 |
| Qwen3.8-Flash-Next (24,576 gate_up expert weights) | 111 s | **68 s
(1.6x)** | 30:23 -> 29:28 |

Qwen's smaller expert weights leave more of each encode to Python launch
work, which the threads share under the GIL. Qwen's KL is identical
token by token (0.287096). GLM's moves within its run-to-run noise
(0.7221 -> 0.7209, paired -0.0011 ± 0.0023, p 0.32), because its
forward, and with it the statistics, varies between runs.

- New 2-GPU test `test_imatrix_encodes_each_gpus_weights_at_once`
(`tests/gpu/torch/quantization/test_ggml_imatrix_gpu.py`). It quantizes
four layers split over two GPUs and the same layers on one GPU. Their
weights must be bit-identical, and encodes on the two GPUs must overlap
in time. It fails when the encodes run in turn.
- `tests/unit/torch/quantization/{test_ggml_imatrix,test_ggml_gptq}.py`
and `plugins/test_fused_experts.py`: 152 passed.
- CI's GPU runner has one GPU, so it skips the new test; I ran it on 2
local GPUs. Not run: offloaded or FSDP2 imatrix with this change, whose
weights take the unchanged sequential path.

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ✅ (same codes; no API change)
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
N/A (performance only)
- Did you get Claude approval on this PR?: ❌

🤖 Generated with [Claude Code](https://claude.com/claude-code)


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
<!-- coderabbit-ai-generated:v1 -->

## Summary by CodeRabbit

* **Performance**
* GGML imatrix calibration can encode weights across multiple GPUs
concurrently, which may reduce calibration time on multi-GPU setups.
* Calibration continues to process weights sequentially when fewer than
two GPUs are available.
* Weights that require special access during calibration are still
processed within that access window.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

#3003)

### What does this PR do?

Type of change: documentation

Transformers runs the linear-attention layers of Qwen3-Next, Qwen3.5,
Qwen3.8, GLM-5.x Flash, Kimi Linear and OLMo Hybrid through readable
PyTorch reference kernels unless the
[flash-linear-attention](https://github.com/fla-org/flash-linear-attention)
kernels are importable. The reference path is slow and needs far more
memory per sample, which caps the layerwise calibration batch. This adds
a note to the `hf_ptq` README and one line to the ptq skill on
installing the kernels: `pip install fla-core`.

- **`fla-core` alone is enough.** The kernels Transformers looks for
(`fla.ops.kda.chunk_kda`,
`fla.ops.gated_delta_rule.chunk_gated_delta_rule`) live in `fla-core`,
which depends only on `einops`. The `flash-linear-attention` package
holds layers and models Transformers does not use, and pins
`transformers>=5.6.2,<6`.
- **GPU only.** Once the kernels are importable, Transformers uses them
for every device, and a CPU forward of these models then fails
(`RuntimeError: 0 active drivers`). That is also why the package is
documented rather than added to `examples/hf_ptq/requirements.txt`,
which the example CI installs.

### Usage

```bash
pip install fla-core
```

### Testing

Layerwise GPTQ on the routed experts at IQ2_XXS on 4x GB300 (oci-jhb),
1,024 x 512 calibration tokens, with #2945 and #3002,
`flash-linear-attention` and `fla-core` 0.6.0:

| | Without the kernels | With them |
|---|---|---|
| GLM-5.3-Flash: batch the probe picks | 16 | 64 (linear-attention
layers fit 181-295 samples instead of 16) |
| GLM-5.3-Flash: layerwise calibration | 35:41 | **16:30** (31:15 at
batch 16) |
| Qwen3.8-Flash-Next: layerwise calibration (batch 64) | 18:23 |
**15:26** |

- With `flash-linear-attention` installed but not `fla-core`, nothing
changed: Transformers could not import `fla.ops.kda.chunk_kda` and kept
the reference path. With `fla-core` alone, it binds
`fla.ops.kda.chunk.chunk_kda` for GLM-5.x and
`fla.ops.gated_delta_rule.chunk.chunk_gated_delta_rule` for Qwen, and a
tiny Qwen3-Next runs on GPU through them; on CPU it fails as described
above.
- The kernels compute slightly differently from the reference: the
unquantized GLM-5.3-Flash under them is 0.018 KL from the reference
logits, which were computed without them. KL to that reference after
GPTQ: Qwen 0.2493 -> 0.2480 (not significant); GLM 0.6189 -> 0.6230 at
batch 16 (+0.0041 +- 0.0045, p 0.36) and 0.6306 at batch 64 (+0.0118 +-
0.0050, p 0.02), paired per chunk over 100 x 512 wikitext-2 tokens.

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ✅
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A (documents
an optional third-party package; nothing added to requirements)
- Did you write any new necessary tests?: N/A
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
N/A
- Did you get Claude approval on this PR?: ❌

🤖 Generated with [Claude Code](https://claude.com/claude-code)


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
<!-- coderabbit-ai-generated:v1 -->

## Summary by CodeRabbit

* **Documentation**
* Added guidance to install `fla-core>=0.5.2` before GPU calibration for
supported models with linear-attention layers, where optimized kernels
can reduce calibration time.
* Clarified that these kernels also run on CPU tensors but can cause CPU
runs to fail, so installation is recommended only in GPU environments.
* Included an example reporting over 2× faster layerwise GPTQ
calibration for GLM-5.3-Flash on 4× GB300.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
Signed-off-by: Chenjie Luo <108829653+cjluo-nv@users.noreply.github.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>